#!/usr/bin/env julia
"""
Quick test of PCG minimization with real cost function
"""

using GSICoreAnalysis
using LinearAlgebra
using Printf
using Random

Random.seed!(42)

println("Quick PCG Minimization Test")
println("="^60)

# Create small configuration
config = AnalysisConfig(
    grid_size = (5, 5, 3),
    ensemble_size = 0,
    use_hybrid = false,
    max_iterations = 10,
    convergence_tol = 1e-6
)

# Create state
state_vector = GSICoreAnalysis.StateVectors.StateVector(config)
state_array = GSICoreAnalysis.CostFunctions.state_vector_to_array(state_vector)
n_state = length(state_array)

println("State size: $n_state")

# Create observations
n_obs = 20
obs_indices = GSICoreAnalysis.CostFunctions.create_observation_indices(n_state, n_obs, :uniform)

# Setup
background = randn(n_state)
observations = randn(n_obs)
obs_errors = ones(n_obs)

# Create cost function
cost_func = GSICoreAnalysis.CostFunctions.CostFunction(config)
cost_func.background_state = background
cost_func.observations = observations
cost_func.observation_errors = obs_errors
cost_func.observation_indices = obs_indices

H_forward, H_adjoint = GSICoreAnalysis.CostFunctions.create_sampling_operator(obs_indices, n_state)
cost_func.observation_operator = H_forward
cost_func.observation_operator_adjoint = H_adjoint

bg_variances = ones(n_state)
obs_variances = obs_errors.^2
GSICoreAnalysis.CostFunctions.setup_diagonal_covariances!(cost_func, bg_variances, obs_variances)

println("\nInitial cost evaluation:")
J0 = GSICoreAnalysis.CostFunctions.evaluate_cost(cost_func, background)
println(@sprintf("  Cost: %.6e", J0))

grad0 = GSICoreAnalysis.CostFunctions.compute_gradient(cost_func, background)
println(@sprintf("  Gradient norm: %.6e", norm(grad0)))

# Create PCG solver with limited iterations
println("\nRunning PCG with max 5 iterations...")
solver = GSICoreAnalysis.Minimization.PCGSolver(config, max_iterations=5)

# Create initial control vector
initial_control = GSICoreAnalysis.ControlVectors.ControlVector(config)

# Run minimization
try
    result = GSICoreAnalysis.Minimization.minimize_cost_function(cost_func, initial_control, solver)

    println("\nResults:")
    println(@sprintf("  Iterations: %d", result.iterations))
    println(@sprintf("  Final cost: %.6e", result.final_cost))
    println(@sprintf("  Final gradient norm: %.6e", result.final_gradient_norm))

    if length(result.cost_history) > 1
        println("\nCost history:")
        for (i, cost) in enumerate(result.cost_history)
            println(@sprintf("  Iter %d: %.6e", i-1, cost))
        end

        reduction = 100.0 * (result.cost_history[1] - result.final_cost) / result.cost_history[1]
        println(@sprintf("\nCost reduction: %.2f%%", reduction))

        if result.final_cost < result.cost_history[1]
            println("✓ Cost decreased successfully")
        else
            println("✗ WARNING: Cost did not decrease")
        end
    end

    println("\n" * "="^60)
    println("Test completed successfully!")
catch e
    println("\nERROR during minimization:")
    println(e)
    println("\nStacktrace:")
    for (exc, bt) in Base.catch_stack()
        showerror(stdout, exc, bt)
        println()
    end
end
